#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
使用通义千问多模态模型将中文JSON键名映射为英文键名
此脚本读取含有中文键名的JSON文件，使用预定义的映射关系替换为英文键名
"""

import os
import json
import re
import argparse
import sys
import time
from typing import Dict, Any, Union, List, Tuple
from qwen_multimodal import query_model

# 配置是否显示详细日志
VERBOSE = True

def log_info(message):
    """打印信息日志"""
    if VERBOSE:
        print(f"[INFO] {message}")

def log_warning(message):
    """打印警告日志"""
    if VERBOSE:
        print(f"[WARNING] {message}")

def log_error(message):
    """打印错误日志"""
    print(f"[ERROR] {message}")

def load_json_file(file_path: str) -> Any:
    """
    加载JSON文件
    
    Args:
        file_path: JSON文件路径
        
    Returns:
        加载的JSON数据
    """
    if not os.path.exists(file_path):
        log_error(f"File not found: {file_path}")
        return None
        
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except json.JSONDecodeError as e:
        log_error(f"Invalid JSON in file {file_path}: {e}")
        return None
    except Exception as e:
        log_error(f"Error reading JSON file {file_path}: {e}")
        return None

def save_json_file(data: Any, file_path: str) -> bool:
    """
    保存JSON数据到文件
    
    Args:
        data: 要保存的JSON数据
        file_path: 输出文件路径
        
    Returns:
        保存是否成功
    """
    try:
        # Create directory only if path contains a directory component
        dir_name = os.path.dirname(file_path)
        if dir_name:
            os.makedirs(dir_name, exist_ok=True)
            
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        log_info(f"Successfully saved to {file_path}")
        return True
    except Exception as e:
        log_error(f"Error saving JSON file {file_path}: {e}")
        return False

def find_mapping_for_key(key: str, mapping: Dict[str, str]) -> str:
    """
    在映射字典中查找键名的映射，确保每个键都能找到最接近的映射
    
    Args:
        key: 要查找的键名
        mapping: 映射字典
        
    Returns:
        映射后的键名，如果没找到则尝试通过语义找到最接近的
    """
    # 直接映射
    if key in mapping:
        log_info(f"为 '{key}' 找到直接映射: '{mapping[key]}'")
        return mapping[key]
    
    # 处理可能的变体（如带括号的版本）
    key_normalized = re.sub(r'[（）\(\)]', '', key)
    if key_normalized in mapping:
        log_info(f"为 '{key}' 找到规范化映射: '{mapping[key_normalized]}'")
        return mapping[key_normalized]
    
    log_info(f"开始为 '{key}' 查找语义映射...")
    
    # 使用通义千问模型进行语义匹配
    prompt = f"""我有一个中文键名需要找到最接近的英文键名映射。请选择语义最接近的一个，即使没有完全匹配也必须返回一个最接近的映射。
中文键名: "{key}"

请从以下可用的映射关系中选择语义最接近的一个:
{json.dumps(mapping, ensure_ascii=False, indent=2)}

只需返回对应的英文键名（值），无需返回中文键。如果多个都可能匹配，选择最合适的一个。
不要返回"NO_MATCH"或类似的词。必须从以上值中选择一个最接近的。"""

    try:
        result = query_model(prompt, preserve_numbers=True).strip()
        
        # 检查返回结果是否在映射值中
        if result in mapping.values():
            log_info(f"为 '{key}' 找到语义匹配: '{result}'")
            return result
            
        # 检查返回值是否包含某个映射值（可能模型返回了额外文字）
        for eng_key in mapping.values():
            if eng_key.lower() in result.lower():
                log_info(f"为 '{key}' 找到部分匹配: '{eng_key}'")
                return eng_key
        
        # 如果模型返回了映射中的键而不是值，尝试找到对应的值
        if result in mapping:
            log_info(f"模型返回了中文键 '{result}'，使用其映射值: '{mapping[result]}'")
            return mapping[result]
        
        # 如果仍然没有找到，使用一个更强制的提示
        log_warning(f"为 '{key}' 第一次匹配失败，尝试强制匹配...")
        
        # 第二次尝试，更明确要求返回映射值
        second_prompt = f"""必须为以下中文键名从给定列表中选择一个最接近的英文映射名称。
中文键名: "{key}"

可选的英文映射值（必须只从这些值中选择）:
{json.dumps(list(mapping.values()), ensure_ascii=False, indent=2)}

只返回列表中最匹配的一个英文名称，不要添加任何额外文字。"""

        result = query_model(second_prompt, preserve_numbers=True).strip()
        
        if result in mapping.values():
            log_info(f"为 '{key}' 第二次尝试找到匹配: '{result}'")
            return result
            
        # 如果仍然失败，生成一个基于中文键的英文名称
        log_warning(f"为 '{key}' 无法找到匹配，生成自定义映射...")
        
        # 使用模型生成合适的英文键名
        gen_prompt = f"""请将这个中文JSON键名 "{key}" 翻译成符合英文命名习惯的下划线分隔的小写英文字段名(snake_case)。
只返回翻译后的英文字段名，不要包含任何其他文字或解释。例如:
- "购买方名称" 应返回 "buyer_name"
- "发票号码" 应返回 "invoice_number"
- "合计金额" 应返回 "total_amount"
"""
        
        custom_key = query_model(gen_prompt, preserve_numbers=True).strip()
        
        # 清理可能的引号和多余空格
        custom_key = custom_key.strip('"\'').strip()
        
        # 确保生成的键是有效的snake_case
        if not custom_key or not all(c.isalnum() or c == '_' for c in custom_key):
            # 回退到简单转换
            log_warning(f"生成的键名 '{custom_key}' 无效，回退到简单转换")
            snake_case = key.lower().replace(' ', '_')
            for char in "（）()[]【】{}。，、；：''""\"":
                snake_case = snake_case.replace(char, '')
            return snake_case
        
        log_info(f"为 '{key}' 生成自定义映射: '{custom_key}'")
        return custom_key
                
    except Exception as e:
        log_error(f"Error using Qwen model for key '{key}': {e}")
        # 回退到简单转换
        snake_case = key.lower().replace(' ', '_')
        for char in "（）()[]【】{}。，、；：''""\"":
            snake_case = snake_case.replace(char, '')
        return snake_case

# 存储所有键映射的字典，用于展示转换结果
key_mapping_results = {}

def split_composite_key(key: str, value: str, mapping: Dict[str, str]) -> List[Tuple[str, str]]:
    """
    拆分复合键名和值到多个字段
    
    Args:
        key: 复合键名
        value: 键值
        mapping: 映射字典
        
    Returns:
        拆分后的(键名,键值)列表
    """
    # 1. 处理预定义的特殊复合字段
    if key == "购买方地址电话" and isinstance(value, str):
        # 使用通义千问模型来拆分地址和电话
        prompt = f"""请将以下文本拆分为地址和电话两部分，格式为JSON:
"{value}"

输出格式：
{{
  "地址": "地址部分",
  "电话": "电话部分"
}}"""
        try:
            result = query_model(prompt, preserve_numbers=True)
            # 尝试解析JSON结果
            split_result = json.loads(result)
            address = split_result.get("地址", "")
            phone = split_result.get("电话", "")
            
            log_info(f"拆分 '{key}' 为地址: '{address}' 和电话: '{phone}'")
            key_mapping_results[key] = ["buyer_address", "buyer_phone"]
            
            return [
                ("buyer_address", address),
                ("buyer_phone", phone)
            ]
        except Exception as e:
            log_error(f"Error splitting address and phone: {e}")
            # 回退到简单分割
            parts = value.split(" ", 1)
            if len(parts) > 1:
                log_info(f"使用简单分割拆分 '{key}'")
                key_mapping_results[key] = ["buyer_address", "buyer_phone"]
                return [
                    ("buyer_address", parts[0]),
                    ("buyer_phone", parts[1])
                ]
            mapped_key = find_mapping_for_key(key, mapping)
            key_mapping_results[key] = [mapped_key]
            return [(mapped_key, value)]
            
    elif key == "销售方地址电话" and isinstance(value, str):
        # 与上面相似的处理方式
        prompt = f"""请将以下文本拆分为地址和电话两部分，格式为JSON:
"{value}"

输出格式：
{{
  "地址": "地址部分",
  "电话": "电话部分"
}}"""
        try:
            result = query_model(prompt, preserve_numbers=True)
            split_result = json.loads(result)
            address = split_result.get("地址", "")
            phone = split_result.get("电话", "")
            
            log_info(f"拆分 '{key}' 为地址: '{address}' 和电话: '{phone}'")
            key_mapping_results[key] = ["seller_address", "seller_phone"]
            
            return [
                ("seller_address", address),
                ("seller_phone", phone)
            ]
        except Exception as e:
            log_error(f"Error splitting address and phone: {e}")
            parts = value.split(" ", 1)
            if len(parts) > 1:
                log_info(f"使用简单分割拆分 '{key}'")
                key_mapping_results[key] = ["seller_address", "seller_phone"]
                return [
                    ("seller_address", parts[0]),
                    ("seller_phone", parts[1])
                ]
            mapped_key = find_mapping_for_key(key, mapping)
            key_mapping_results[key] = [mapped_key]
            return [(mapped_key, value)]
            
    elif key == "购买方开户行及账号" and isinstance(value, str):
        # 拆分开户行和账号
        prompt = f"""请将以下文本拆分为银行和账号两部分，格式为JSON:
"{value}"

输出格式：
{{
  "银行": "银行名称",
  "账号": "账号"
}}"""
        try:
            result = query_model(prompt, preserve_numbers=True)
            split_result = json.loads(result)
            bank = split_result.get("银行", "")
            account = split_result.get("账号", "")
            
            log_info(f"拆分 '{key}' 为银行: '{bank}' 和账号: '{account}'")
            key_mapping_results[key] = ["buyer_bank", "buyer_bank_account"]
            
            return [
                ("buyer_bank", bank),
                ("buyer_bank_account", account)
            ]
        except Exception as e:
            log_error(f"Error splitting bank and account: {e}")
            mapped_key = find_mapping_for_key(key, mapping)
            key_mapping_results[key] = [mapped_key]
            return [(mapped_key, value)]
            
    elif key == "销售方开户行及账号" and isinstance(value, str):
        # 拆分开户行和账号
        prompt = f"""请将以下文本拆分为银行和账号两部分，格式为JSON:
"{value}"

输出格式：
{{
  "银行": "银行名称",
  "账号": "账号"
}}"""
        try:
            result = query_model(prompt, preserve_numbers=True)
            split_result = json.loads(result)
            bank = split_result.get("银行", "")
            account = split_result.get("账号", "")
            
            log_info(f"拆分 '{key}' 为银行: '{bank}' 和账号: '{account}'")
            key_mapping_results[key] = ["seller_bank", "seller_bank_account"]
            
            return [
                ("seller_bank", bank),
                ("seller_bank_account", account)
            ]
        except Exception as e:
            log_error(f"Error splitting bank and account: {e}")
            mapped_key = find_mapping_for_key(key, mapping)
            key_mapping_results[key] = [mapped_key]
            return [(mapped_key, value)]
            
    elif key.startswith("价税合计") and isinstance(value, str):
        if "大写" in key:
            log_info(f"将 '{key}' 映射为 'total_amount_in_words'")
            key_mapping_results[key] = ["total_amount_in_words"]
            return [("total_amount_in_words", value)]
        elif "小写" in key:
            log_info(f"将 '{key}' 映射为 'total_amount_in_digits'")
            key_mapping_results[key] = ["total_amount_in_digits"]
            return [("total_amount_in_digits", value)]
        else:
            mapped_key = find_mapping_for_key(key, mapping)
            key_mapping_results[key] = [mapped_key]
            return [(mapped_key, value)]
    
    # 2. 检测并处理其他可能的复合键
    # 如果键名中包含"及"、"与"、"和"等连接词，可能是复合键
    if len(key) > 3 and isinstance(value, str) and any(word in key for word in ["及", "与", "和"]):
        # 使用通义千问模型识别此键是否可拆分，以及如何拆分
        prompt = f"""请分析这个中文键名 "{key}" 是否是一个复合字段（包含多个信息的字段）。
如果是复合字段，请将其拆分成单独的字段，并给出每个子字段的英文名称(snake_case格式)和对应的值。
原始值: "{value}"

输出格式JSON:
{{
  "is_composite": true或false,
  "fields": [
    {{"name": "英文字段名1", "value": "对应值1"}},
    {{"name": "英文字段名2", "value": "对应值2"}}
  ]
}}

如果不是复合字段，则返回is_composite为false。"""

        try:
            result = query_model(prompt, preserve_numbers=True)
            analysis = json.loads(result)
            
            if analysis.get("is_composite", False) and isinstance(analysis.get("fields"), list):
                fields = analysis.get("fields", [])
                if len(fields) > 1:
                    log_info(f"检测到复合键 '{key}'，拆分为 {len(fields)} 个子字段")
                    field_names = [field.get("name") for field in fields]
                    key_mapping_results[key] = field_names
                    return [(field.get("name"), field.get("value")) for field in fields]
        except Exception as e:
            log_error(f"Error analyzing composite key '{key}': {e}")
    
    # 3. 查找可能的地址+电话组合
    if isinstance(value, str) and len(value) > 10 and (" " in value or "　" in value):
        # 检查值是否可能是地址+电话组合
        has_number = bool(re.search(r'\d{5,}', value))
        if has_number:
            prompt = f"""请检测以下文本是否包含地址和电话号码。如果包含，请拆分成地址和电话两部分，格式为JSON:
"{value}"

输出格式：
{{
  "contains_phone": true或false,
  "地址": "地址部分",
  "电话": "电话部分"
}}"""
            try:
                result = query_model(prompt, preserve_numbers=True)
                split_result = json.loads(result)
                
                if split_result.get("contains_phone", False):
                    address = split_result.get("地址", "")
                    phone = split_result.get("电话", "")
                    
                    prefix = ""
                    if "购买" in key or "买方" in key:
                        prefix = "buyer_"
                    elif "销售" in key or "卖方" in key:
                        prefix = "seller_"
                    
                    address_key = f"{prefix}address"
                    phone_key = f"{prefix}phone"
                    
                    log_info(f"检测到可能的地址电话组合 '{key}'，拆分为 {address_key} 和 {phone_key}")
                    key_mapping_results[key] = [address_key, phone_key]
                    
                    return [
                        (address_key, address),
                        (phone_key, phone)
                    ]
            except Exception as e:
                log_error(f"Error detecting address and phone: {e}")
            
    # 默认情况：使用映射找到对应的英文键名
    mapped_key = find_mapping_for_key(key, mapping)
    key_mapping_results[key] = [mapped_key]
    return [(mapped_key, value)]

def translate_json_keys(json_data: Union[Dict, List], key_mapping: Dict[str, str]) -> Union[Dict, List, Any]:
    """
    递归翻译JSON结构中的所有中文键名为英文
    
    Args:
        json_data: 要翻译的JSON数据
        key_mapping: 键名映射字典
        
    Returns:
        翻译后的JSON数据
    """
    if isinstance(json_data, dict):
        new_dict = {}
        
        # 特殊处理商品列表
        if "商品列表" in json_data:
            items_key = find_mapping_for_key("商品列表", key_mapping)
            if isinstance(json_data["商品列表"], list):
                log_info(f"处理商品列表字段，映射为 '{items_key}'")
                new_dict[items_key] = translate_json_keys(json_data["商品列表"], key_mapping)
        
        # 处理其他所有非列表项
        for key, value in json_data.items():
            if key != "商品列表":  # 已经处理过商品列表
                if isinstance(value, (dict, list)):
                    # 递归处理嵌套结构
                    new_key = find_mapping_for_key(key, key_mapping)
                    log_info(f"处理嵌套结构 '{key}' -> '{new_key}'")
                    new_dict[new_key] = translate_json_keys(value, key_mapping)
                else:
                    # 处理复合键
                    key_value_pairs = split_composite_key(key, value, key_mapping)
                    for new_key, new_value in key_value_pairs:
                        new_dict[new_key] = new_value
        
        return new_dict
    elif isinstance(json_data, list):
        # 对于列表，翻译每个项目（如果是字典或列表）
        log_info(f"处理列表类型，包含 {len(json_data)} 个项目")
        return [translate_json_keys(item, key_mapping) if isinstance(item, (dict, list)) else item for item in json_data]
    else:
        # 对于基本类型（字符串、数字等），原样返回
        return json_data

def map_chinese_to_english(input_json_path: str, mapping_json_path: str, output_json_path: str = None) -> Dict:
    """
    将中文JSON键名映射为英文键名
    
    Args:
        input_json_path: 输入JSON文件路径
        mapping_json_path: 键名映射文件路径
        output_json_path: 可选的输出文件路径
        
    Returns:
        翻译后的JSON数据
    """
    # 清空映射结果记录
    global key_mapping_results
    key_mapping_results = {}
    
    # 加载输入JSON文件
    data = load_json_file(input_json_path)
    if data is None:
        return None
        
    # 加载映射JSON文件
    key_mapping = load_json_file(mapping_json_path)
    if key_mapping is None:
        return None
    
    log_info(f"已加载映射文件，包含 {len(key_mapping)} 个映射关系")
    
    # 翻译JSON数据
    start_time = time.time()
    log_info("开始处理JSON数据...")
    translated_data = translate_json_keys(data, key_mapping)
    log_info(f"处理完成，耗时 {time.time() - start_time:.2f} 秒")
    
    # 保存键名映射结果
    if output_json_path:
        mapping_result_path = f"{os.path.splitext(output_json_path)[0]}_mapping_result.json"
        log_info(f"保存键名映射结果到 {mapping_result_path}")
        save_json_file(key_mapping_results, mapping_result_path)
    
    # 保存翻译后的JSON（如果提供了输出路径）
    if output_json_path:
        save_json_file(translated_data, output_json_path)
    
    # 打印映射摘要
    log_info("\n映射摘要:")
    for key, mapped_keys in key_mapping_results.items():
        if len(mapped_keys) == 1:
            log_info(f"  '{key}' -> '{mapped_keys[0]}'")
        else:
            log_info(f"  '{key}' -> {mapped_keys}")
    
    return translated_data

def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description='将中文JSON键名映射为英文键名')
    parser.add_argument('input_file', help='输入JSON文件路径')
    parser.add_argument('mapping_file', help='键名映射文件路径')
    parser.add_argument('-o', '--output', help='输出JSON文件路径')
    parser.add_argument('-v', '--verbose', action='store_true', help='启用详细日志输出')
    parser.add_argument('-q', '--quiet', action='store_true', help='禁用详细日志输出')
    
    # 如果没有参数，显示帮助
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)
        
    args = parser.parse_args()
    
    # 设置日志级别
    global VERBOSE
    if args.verbose:
        VERBOSE = True
    if args.quiet:
        VERBOSE = False
    
    # 检查文件是否存在
    if not os.path.exists(args.input_file):
        log_error(f"错误: 输入文件不存在: {args.input_file}")
        sys.exit(1)
        
    if not os.path.exists(args.mapping_file):
        log_error(f"错误: 映射文件不存在: {args.mapping_file}")
        sys.exit(1)
    
    # 默认输出路径（如果未指定）
    output_file = args.output or f"{os.path.splitext(args.input_file)[0]}_english.json"
    
    # 执行映射
    result = map_chinese_to_english(args.input_file, args.mapping_file, output_file)
    
    if result:
        print("映射成功完成!")
    else:
        log_error("映射过程中发生错误。")
        sys.exit(1)

if __name__ == "__main__":
    main() 